﻿using Dou.GPU;
using Unity.Burst;
using Unity.Collections;
using Unity.Collections.LowLevel.Unsafe;
using Unity.Jobs;
using Unity.Mathematics;
using UnityEngine;
using Random = Unity.Mathematics.Random;

namespace Simple.GPUInstancing.SampleTest
{
    /// <summary>
    /// 创建移动小方块并渲染
    /// </summary>
    public unsafe class Sample1Script : MonoBehaviour
    {
        public Mesh mesh;
        public Material material;
        public int count = 100000;

        private int _batchId;
        private static readonly int ID_TransformBuf = Shader.PropertyToID("transformBuf");

        private NativeArray<SampleData> _dataArr;
        
        void Start()
        {
            // 注册批次
            _batchId = GPUInstancingMgr.instance.RegisterBatch(mesh, material);
            // 注册材质属性
            GPUInstancingMgr.instance.RegisterProperty(_batchId, ID_TransformBuf, GPUDefine.Stride_Float4x4);

            InitData();
        }

        private void InitData()
        {
            _dataArr = new NativeArray<SampleData>(count, Allocator.Persistent);
            
            var job = new InitJob();
            job.array = _dataArr;
            var jobHandle = job.Schedule(count, 64);
            jobHandle.Complete();
        }

        [BurstCompile]
        public struct InitJob : IJobParallelFor
        {
            public NativeArray<SampleData> array;
            
            public void Execute(int index)
            {
                var random = Random.CreateFromIndex((uint) index);

                var ptr = (SampleData*) array.GetUnsafePtr();
                
                var data = ptr + index;
                
                data->pos = new float3(random.NextFloat(-50f, 50f), random.NextFloat(-50f, 50f), random.NextFloat(-50f, 50f));
                data->rotation = quaternion.Euler(random.NextFloat(360f), random.NextFloat(360f), random.NextFloat(360f));
                var scale = random.NextFloat(0.5f, 1.5f);
                data->scale = new float3(scale, scale, scale);

                data->forward = math.normalize(math.forward(data->rotation));
                data->rotateAdd = new float3(random.NextFloat(10f), random.NextFloat(10f), random.NextFloat(10f));
            }
        }

        void Update()
        {
            UpdateData();
        }

        private void UpdateData()
        {
            var transArr = new NativeArray<float4x4>(count, Allocator.TempJob);
            
            var job = new UpdateJob();
            job.dataArr = _dataArr;
            job.transArr = transArr;
            job.deltaTime = Time.deltaTime;
            var jobHandle = job.Schedule(count, 64);
            jobHandle.Complete();
            
            // 设置数据
            GPUInstancingMgr.instance.SetFloat4x4(_batchId, ID_TransformBuf, ref transArr);
            // 绘制全部数据
            GPUInstancingMgr.instance.Draw();

            transArr.Dispose();
        }

        [BurstCompile]
        public struct UpdateJob : IJobParallelFor
        {
            public NativeArray<SampleData> dataArr;
            public NativeArray<float4x4> transArr;
            public float deltaTime;
            
            public void Execute(int index)
            {
                var dataPtr = (SampleData*) dataArr.GetUnsafePtr();
                var transPtr = (float4x4*) transArr.GetUnsafePtr();
                
                var data = dataPtr + index;

                data->pos += data->forward * deltaTime;
                data->rotation = math.mul(data->rotation, quaternion.Euler(data->rotateAdd * deltaTime));
                
                var trans = float4x4.TRS(data->pos, data->rotation, data->scale);
                *(transPtr + index) = trans;
            }
        }

        private void OnDestroy()
        {
            _dataArr.Dispose();
        }

        public struct SampleData
        {
            public float3 pos;
            public quaternion rotation;
            public float3 scale;

            public float3 forward; // 1秒向前移动的位移
            public float3 rotateAdd; // 1秒旋转的增量
        }
    }
}
